# -*- coding:utf-8 -*-

# @Time  : 2024/1/17 3:42 PM
# @Author: chenyong
# @Email : chenyong@lingxi.ai
# @File  : generate_plaaner_agent_data.py

import pandas as pd
import asyncio
from bot.insurance_planner_gpt.planner import PlannerChat
from bot.insurance_planner_gpt.context_local_util import context_data
# from bot.insurance_planner_gpt.agent.user import User
import uuid
from data_generate import utils
import traceback
import random
import os
import openpyxl

def read_excel(excel_path):
    conversations = {}
    df = pd.read_excel(excel_path)
    for i in range(len(df)):
        id = df.loc[i, "对话id"]
        role = df.loc[i, "角色"]
        content = df.loc[i, "对话内容"]
        if not id:
            continue
        if id not in conversations:
            conversations[id] = []
        conversations[id].append({"role": role, "content": content})

    return conversations


async def mock_bot_dialogue(path):
    train_data = []
    if os.path.exists(path):
        train_data = utils.jload(path)

    conversations = read_excel("../xiaozhu/conversation_20240111.xlsx")

    nums = 0
    for _id, _conversation in conversations.items():
        nums += 1
        if nums == 3:
            break
        try:
            bot = PlannerChat(None)
            session_id = _id + "_new"
            context_data.set({"session_id": session_id, "message_id": 'message_id'})
            context = []
            instruct_dict_new = {}
            instruct_dict_new["id"] = session_id
            instruct_dict_new["messages"] = []
            i = 0
            while i < len(_conversation) - 1:
                j = i+1
                _role = _conversation[i]['role']
                _content = _conversation[i]['content']
                conversation_dict = {}
                conversation_dict["role"] = _role
                conversation_dict["content"] = _content
                instruct_dict_new["messages"].append(conversation_dict)
                context.append(conversation_dict)

                _role = _conversation[j]['role']
                _content = _conversation[j]['content']
                conversation_dict = {}
                conversation_dict["role"] = _role
                conversation_dict["content"] = _content
                instruct_dict_new["messages"].append(conversation_dict)
                context.append(conversation_dict)

                res, other_messages = await bot.async_reply(context, session_id, None)

                conversation_dict = {}
                conversation_dict["role"] = "assistant"
                conversation_dict["content"] = res
                instruct_dict_new["messages"].append(conversation_dict)
                # context.append(conversation_dict)

                i = j + 1

            if i == len(_conversation) - 1:
                _role = _conversation[i]['role']
                _content = _conversation[i]['content']
                conversation_dict = {}
                conversation_dict["role"] = _role
                conversation_dict["content"] = _content
                instruct_dict_new["messages"].append(conversation_dict)
                context.append(conversation_dict)

            train_data.append(instruct_dict_new)
            utils.jdump(train_data, path)
        except Exception as e:
            traceback.print_exc()
            # 打印堆栈
            print(e)


def add_empty_row(group):
    return pd.concat([group, pd.DataFrame([[''] * len(group.columns)], columns=group.columns)],
                     ignore_index=True)


def conversation2csv(file_path):
    datas = utils.jload(file_path)
    df = pd.DataFrame(columns=['角色', '内容'])

    for data in datas:
        messages = data['messages']
        for text in messages:
            df = df.append(
                {"角色": text['role'],
                 "内容": text['content']},
                ignore_index=True)
        df = df.reset_index(drop=True).append(
            {"角色": "",
             "内容": ""},
            ignore_index=True)
    df.to_csv("gen_conversation.csv", index=False, encoding='utf-8_sig')


def read_agent_info(file_path):
    datas = utils.jload(file_path)
    resp = {}
    for data in datas:
        _id = data["id"]
        _time = data["time"]
        _task = data["task"]
        _model_name = data["model_name"]
        _conversations = data["conversations"]
        for conv in _conversations:
            if conv["from"] == 'gpt':
                value = conv["value"]
        if _id not in resp:
            resp[_id] = {}
        if _task not in resp[_id]:
            resp[_id][_task] = {}
        resp[_id][_task]["model_name"] = _model_name
        resp[_id][_task]["time"] = _time
        resp[_id][_task]["value"] = value

    return resp


def conversation2xlsx(dialogue_path, agent_path):
    agent_res = read_agent_info(agent_path)
    datas = utils.jload(dialogue_path)
    ids = []
    roles = []
    contents = []
    model_names = []
    times = []
    # df = pd.DataFrame(columns=['对话id', '角色', '对话内容', '对话时间', '模型名称'])
    for data in datas:
        _id = data["id"]
        _messages = data["messages"]
        for text in _messages:
            role = text['role']
            content = text['content']
            if "message_id" in text:
                key = _id + ":" + text["message_id"]
                values = agent_res[key]
                for k, v in values.items():
                    model_name = v["model_name"]
                    time = v["time"]
                    value = v["value"]
                    title = ""
                    if k == "QuestioningDisputeResolution":
                        title = "用户问题"
                    if k == "UserInfoExtract":
                        title = "用户信息"
                    if k == "ProductCausalSolution":
                        title = "因果推断"
                    if k == "PlanExplain":
                        title = "方案对比"
                    if k == "PlanAssume":
                        title = "因果图"
                    if title != "":
                        ids.append(_id)
                        roles.append(role)
                        contents.append({"title": title, "content": value})
                        model_names.append(model_name)
                        times.append(time)

                ids.append(_id)
                roles.append(role)
                contents.append(content)
                model_names.append("")
                times.append(time)
            else:
                ids.append(_id)
                roles.append(role)
                contents.append(content)
                model_names.append("")
                times.append("")

        ids.append("")
        roles.append("")
        contents.append("")
        model_names.append("")
        times.append("")

    df = pd.DataFrame({"对话id": ids, "角色": roles, "对话内容": contents, "对话时间": times, "模型名称": model_names})
    df.to_excel("conversation_1115_1.xlsx", index=False)


if __name__ == '__main__':
    import datetime

    now_date = datetime.datetime.now().strftime("%Y-%m-%d")
    file_path = "../data/planner/bot_dialogue/bot_dialogue_" + now_date + ".json"  # 文件夹路径
    asyncio.run(mock_bot_dialogue(file_path))


    # dialogue_path = "../data/planner/bot_dialogue/bot_dialogue_2023-11-15.json"
    # agent_path = "../internal_server/train_data/llm/2023-11-15.json"
    # conversation2xlsx(dialogue_path, agent_path)

